from Utils.util import *

class Prim_Maze:
    """
    随机prim生成迷宫
    """
    def __init__(self, rows, columns) -> None:
        self.num_rows = rows
        self.num_cols = columns

    def run(self) -> list:
        num_rows = self.num_rows
        num_cols = self.num_cols
        
        M = np.zeros((num_rows,num_cols,5), dtype=np.uint8)

        image = np.zeros((num_rows*10,num_cols*10), dtype=np.uint8)
        
        r = 0
        c = 0
        history = [(r,c)]

        while history: 

            r,c = random.choice(history)
            M[r,c,4] = 1
            history.remove((r,c))
            check = []
            if c > 0:
                if M[r,c-1,4] == 1:
                    check.append('L')
                elif M[r,c-1,4] == 0:
                    history.append((r,c-1))
                    M[r,c-1,4] = 2
            if r > 0:
                if M[r-1,c,4] == 1: 
                    check.append('U') 
                elif M[r-1,c,4] == 0:
                    history.append((r-1,c))
                    M[r-1,c,4] = 2
            if c < num_cols-1:
                if M[r,c+1,4] == 1: 
                    check.append('R')
                elif M[r,c+1,4] == 0:
                    history.append((r,c+1))
                    M[r,c+1,4] = 2 
            if r < num_rows-1:
                if M[r+1,c,4] == 1: 
                    check.append('D') 
                elif  M[r+1,c,4] == 0:
                    history.append((r+1,c))
                    M[r+1,c,4] = 2

            if len(check):
                move_direction = random.choice(check)
                if move_direction == 'L':
                    M[r,c,0] = 1
                    c = c-1
                    M[r,c,2] = 1
                if move_direction == 'U':
                    M[r,c,1] = 1
                    r = r-1
                    M[r,c,3] = 1
                if move_direction == 'R':
                    M[r,c,2] = 1
                    c = c+1
                    M[r,c,0] = 1
                if move_direction == 'D':
                    M[r,c,3] = 1
                    r = r+1
                    M[r,c,1] = 1
                    
        Map = draw_maze(num_rows, num_cols, M, image)
        return Map